#include	<x86intrin.h>
#include	<sys/time.h>
#include	<algorithm>
#include	<iostream>
#include	<unistd.h>
#include	"utf16.h"
#include	<cstring>
#include	<sstream>
#include	<cfloat>
#include	<cstdio>
#include	<vector>
#include	<cmath>
#include	<ctime>
#include	<omp.h>
using	namespace	std;
#define Q8B     0.105409255338946f
uint64_t	prng=time(NULL);
static inline uint64_t wyrand(uint64_t *seed){	*seed+=0xa0761d6478bd642full;	uint64_t  see1=*seed^0xe7037ed1a0b428dbull;	see1*=(see1>>32)|(see1<<32);	return  (*seed*((*seed>>32)|(*seed<<32)))^((see1>>32)|(see1<<32));}
static inline double wy2u01(uint64_t r){ const double _wynorm=1.0/(1ull<<52); return (r>>12)*_wynorm;}
static inline float wy2gau(uint64_t r){ const float	_wynorm=1.0/(1ull<<20); return ((r&0x1fffff)+((r>>21)&0x1fffff)+((r>>42)&0x1fffff))*_wynorm-3.0f;}
static inline void _wymum(uint64_t *A,	uint64_t *B){	uint64_t	hh=(*A>>32)*(*B>>32), hl=(*A>>32)*(uint32_t)*B, lh=(uint32_t)*A*(*B>>32), ll=(uint64_t)(uint32_t)*A*(uint32_t)*B;	*A=((hl>>32)|(hl<<32))^hh;	*B=((lh>>32)|(lh<<32))^ll;}
static inline uint64_t	_wyhash64(uint64_t	A,	uint64_t	B){	A^=0xa0761d6478bd642full; B^=0xa0761d6478bd642full;	_wymum(&A,&B);	A^=0xa0761d6478bd642full; B^=0xa0761d6478bd642full;	_wymum(&A,&B);	return	A^B;}
template<uint32_t	N>
struct	Data8{
	char	*data;
	Data8(){	data=(char*)aligned_alloc(64,N);	}
//	Data8(){	data=(char*)_aligned_malloc(N,64);	}
	~Data8(){	free(data);	}
//	~Data8(){	_aligned_free(data);	}
	bool	load(FILE	*F){	return	fread(data,N,1,F)==1;	}
};
template<uint32_t	N>
struct	Data{
	float	*data;
	Data(){	data=(float*)aligned_alloc(64,N*sizeof(float));	}
//	Data(){	data=(float*)_aligned_malloc(N*sizeof(float),64);	}
	~Data(){	free(data);	}
//	~Data(){	_aligned_free(data);	}
	void	load(FILE	*F){
		Data8<N>	tmp;	tmp.load(F);	
		for(uint32_t	i=0;	i<N;	i++)	data[i]=Q8B*tmp.data[i];
	}
};
const   int vnni_const=90;
template<uint32_t   R,	uint32_t   C>
class	vnnia{
public:
	float scale[C]; uint8_t  *data;
	vnnia(){	data=(uint8_t*)aligned_alloc(32,R*C); 	}
	~vnnia(){	free(data);	}
	void	float2vnni(float	*p,	uint32_t	c){
		uint8_t	*q=data+c*R;	float ma=-FLT_MAX;
		for(uint32_t	i=0;	i<R;	i++)	ma=max(ma,fabsf(p[i]));
		scale[c]=vnni_const/ma;
		for(uint32_t	i=0;	i<R;	i++)	q[i]=roundf(p[i]*scale[c]+vnni_const);
	}
	void	char2vnni(char	*p,	uint32_t	c){
		uint8_t	*q=data+c*R;	scale[c]=1;
		for(uint32_t	i=0;	i<R;	i++)	q[i]=p[i]+vnni_const;
	}	
};
template<uint32_t   R,	uint32_t   C>
class	vnnib{
public:
	float scale[C];	int	sum[C];	char  *data;
	vnnib(){	data=(char*)aligned_alloc(32,R*C);	}
	~vnnib(){	free(data);	}
	void	float2vnni(float	*p,	uint32_t	c){
		char  *q=data+c*R;	float ma=0;
		for(uint32_t  i=0;  i<R;  i++)	ma=max(ma,fabsf(p[i]));
		scale[c]=vnni_const/ma;	int	s=0;
		for(uint32_t  i=0;  i<R;  i++){	q[i]=roundf(p[i]*scale[c]);	s+=q[i];	}
		sum[c]=s;
	}
};
template<uint32_t   R,	uint32_t   C,	uint32_t	H>
class	vnnic{
public:
	const	float scale=vnni_const/sqrtf(H);
	uint8_t  *data;
	vnnic(){	data=(uint8_t*)aligned_alloc(32,R*C);	for(uint32_t	i=0;	i<R*C;	i++)	data[i]=vnni_const;	}
	~vnnic(){	free(data);	}
	void	float2vnni(float	*p,	uint32_t	c){
		for(uint32_t	i=0;	i<C;	i++)	data[i*R+c]=roundf(fminf(fmaxf(p[i]*scale+vnni_const,0),2*vnni_const));
	}
};
static	inline	int	_mm256_reduce_add_epi32(__m256i	x){	x=_mm256_hadd_epi32(x,x);	x=_mm256_hadd_epi32(x,x);	int	*p=(int*)&x;	return	p[0]+p[4];	}
template<uint32_t	R,	uint32_t	C0,	uint32_t	C1>
float	dot(vnnia<R,C0>	&ma,	uint32_t	ca,	vnnib<R,C1>	&mb,	uint32_t	cb){
	uint8_t	*p=ma.data+ca*R;	char *q=mb.data+cb*R;
	__m256i	v1,	v2,	v3=_mm256_setzero_si256();
	for(uint32_t  i=0;  i<R;  i+=32){	v1=_mm256_load_si256((__m256i*)(p+i));	v2 =_mm256_load_si256((__m256i*)(q+i));	v3 =_mm256_dpbusds_epi32(v3, v1, v2);	}
	float z=_mm256_reduce_add_epi32(v3);
	return   (z-vnni_const*mb.sum[cb])/(mb.scale[cb]*ma.scale[ca]);
};
template<uint32_t	R,	uint32_t	C0,	uint32_t	C1,	uint32_t	H>
float	dot(vnnic<R,C0,H>	&ma,	uint32_t	ca,	vnnib<R,C1>	&mb,	uint32_t	cb){
	uint8_t	*p=ma.data+ca*R;	char *q=mb.data+cb*R;
	__m256i	v1,	v2,	v3=_mm256_setzero_si256();
	for(uint32_t  i=0;  i<R;  i+=32){	v1=_mm256_load_si256((__m256i*)(p+i));	v2 =_mm256_load_si256((__m256i*)(q+i));	v3 =_mm256_dpbusds_epi32(v3, v1, v2);	}
	float z=_mm256_reduce_add_epi32(v3);
	return   (z-vnni_const*mb.sum[cb])/(mb.scale[cb]*ma.scale);
};
template<uint32_t	R0,	uint32_t	R1>
struct	linear{
	vnnia<R0,R1>	vwei;
	vnnib<R0,1>	vinp;
	Data<R1>	out;
	void	load(FILE	*F){
		Data8<R0*R1>	wei;	wei.load(F);
		for(uint32_t	i=0;	i<R1;	i++)	vwei.char2vnni(wei.data+i*R0,i);
	}
	void	fw(Data<R0*1>	&inp,	uint32_t	col=0){
		float	*ou=out.data+col*R1,	alf=Q8B/sqrtf(R0);
		vinp.float2vnni(inp.data,0);
		#pragma omp parallel for
		for(uint32_t	i=0;	i<R1;	i++)	ou[i]=alf*dot(vwei,i,vinp,0);
	}
};
template<uint32_t	R,	uint32_t	H>
void	layernorm(Data<R>	&inp){
	uint32_t	r=R/H;
	#pragma omp parallel for
	for(uint32_t	h=0;	h<H;	h++){
		float	sum=0,	nor=0,	*p=inp.data+h*r;	
		for(uint32_t	i=0;	i<r;	i++){	sum+=p[i];	nor+=p[i]*p[i];	}
		sum/=r;	nor=fmaxf(nor-sum*sum*r,1e-18f);	nor=sqrtf(r/nor);
		for(uint32_t	i=0;	i<r;	i++)	p[i]=(p[i]-sum)*nor;
	}	
}
template<uint32_t	R>
void	softmax(Data<R>	&inp){
	float	ma=-FLT_MAX,	sum=0;
	for(uint32_t	i=0;	i<R;	i++)	if(inp.data[i]>ma)	ma=inp.data[i];
	for(uint32_t	i=0;	i<R;	i++)	sum+=(inp.data[i]=expf(inp.data[i]-ma));
	for(uint32_t	i=0;	i<R;	i++)	inp.data[i]/=sum;
}
template<uint32_t	R,	uint32_t	C,	uint32_t	H>
struct	sexy{
	vnnib<R/H,1>	vq[H];
	vnnia<R/H,C>	vk[H];
	vnnib<C,1>	va[H];
	vnnic<C,R,R/H>	pn;	
	Data<H*C>	pe;
	linear<R,4*R>	x;
	linear<R,R>	o;
	Data<R>	tmp,&out=o.out;
	void	load(FILE	*F){	pe.load(F);	x.load(F);	o.load(F);	}
	void	fw(Data<R>	&inp,	uint32_t	col,	uint32_t	para){
		x.fw(inp);	layernorm<4*R,4*H>(x.out);	pn.float2vnni(x.out.data+2*R,col);
		#pragma omp parallel for
		for(uint32_t	h=0;	h<H;	h++){
			vk[h].float2vnni(x.out.data+h*(R/H),col);
			vq[h].float2vnni(x.out.data+R+h*(R/H),0);
			float	alf=1/sqrtf(R/H),	a[C]={};

			for(uint32_t	j=para;	j<C;	j++){	uint32_t	i=(j+1+col)%C;	a[i]=expf(dot(vk[h],i,vq[h],0)*alf+pe.data[h*C+C-1-j]);	}
			va[h].float2vnni(a,0);
			for(uint32_t	i=0;	i<R/H;	i++)	tmp.data[h*(R/H)+i]=dot(pn,h*(R/H)+i,va[h],0);
		}
		layernorm<R,H>(tmp);	
		for(uint32_t	i=0;	i<R;	i++)	tmp.data[i]*=x.out.data[3*R+i];
		o.fw(tmp);
		for(uint32_t	i=0;	i<R;	i++)	o.out.data[i]+=inp.data[i];
	}
};
template<uint32_t	R,	uint32_t	C,	uint32_t	H>
struct	self{
	static	Data<R>	tmp;
	linear<R,2*R>	u;
	linear<R,R>	o;
	Data<R>	&out=o.out;
	void	load(FILE* F){	 u.load(F);	o.load(F);	}
	void	fw(Data<R>	&inp){
		u.fw(inp);	layernorm<2*R,2*H>(u.out);
		for(uint32_t	i=0;	i<R;	i++)	tmp.data[i]=u.out.data[i]*u.out.data[R+i];
		o.fw(tmp);
		for(uint32_t	i=0;	i<R;	i++)	o.out.data[i]+=inp.data[i];	
	}
};
template<uint32_t	R,	uint32_t	C,	uint32_t	H>
Data<R>	self<R,C,H>::tmp;
template<uint32_t	R,	uint32_t	C,	uint32_t	H>
struct	wyGPT{
	self<R,C,H>	a;
	sexy<R,C,H>	b;
	self<R,C,H>	c;	
	Data<R>	&out=c.out;
	void	load(FILE* F){	 a.load(F);	b.load(F);	c.load(F);	}
	void	fw(Data<R>	&inp,	uint32_t	col,	uint32_t	para){
		a.fw(inp);
		b.fw(a.out,col,para);
		c.fw(b.out);		
	}
};
template<uint32_t	C,	uint32_t	E,	uint32_t	D,	uint32_t	H,	uint32_t	O>
struct	Denisovan{
	uint32_t	curr=0;
	uint16_t	code[O],table[65536];
	vector<bool>	forbid;
	Data<E>	emb;
	wyGPT<E,C,H>	tra[D];
	linear<E,O>	out;
	float	vs[O];
	Denisovan(){
		forbid.resize(0x10000);
		string	s="　“”《》（）【】「」";
		vector<uint16_t>	v;	v.resize(s.size());
		v.resize(utf8_to_utf16((uint8_t*)s.data(),s.size(),v.data(),v.size()));
		for(uint32_t	i=0;	i<v.size();	i++)	forbid[v[i]]=true;
	}
	bool	load(const	char	*F){
		FILE* f=fopen(F, "rb");
		if(f==NULL)	return	false;
		uint32_t	x;
		if(fread(&x,4,1,f)!=1||x!=C)	fprintf(stderr,"C=%u\n",x);
		if(fread(&x,4,1,f)!=1||x!=E)	fprintf(stderr,"E=%u\n",x);
		if(fread(&x,4,1,f)!=1||x!=D)	fprintf(stderr,"D=%u\n",x);
		if(fread(&x,4,1,f)!=1||x!=H)	fprintf(stderr,"H=%u\n",x);
		if(fread(&x,4,1,f)!=1||x!=O)	fprintf(stderr,"O=%u\n",x);		
		if(fread(code,O*2,1,f)!=1)	return	false;
		for(uint32_t i=0; i<D; i++)	tra[i].load(f);
		out.load(f);	fclose(f);
		for(uint32_t	i=0;	i<65536;	i++)	table[i]=0;
		for(uint32_t	i=0;	i<O;	i++)	table[code[i]]=i;
		return	true;
	}
	uint16_t	sample(uint16_t	*x,	uint16_t	*p){
		uint32_t	para=p+C-1>=x?p+C-1-x:0;
		for(uint32_t	r=0;	r<E;	r++)	emb.data[r]=(_wyhash64(*x,r)&1)*2-1.0f;
		for(uint32_t	d=0;	d<D;	d++)	tra[d].fw(d?tra[d-1].out:emb,curr,para);
		layernorm<E,H>(tra[D-1].out);	out.fw(tra[D-1].out);
		for(uint32_t    i=0;    i<O;    i++)	out.out.data[i]=(code[i]<128||forbid[code[i]])?-1000:M_SQRT2*(out.out.data[i]-vs[i]);
		softmax<O>(out.out);
		double	sum=0;	for(uint32_t    i=0;    i<O;    i++)	sum+=(out.out.data[i]=fmaxf(out.out.data[i]-1.0f/O,0));
		double  ran=wy2u01(wyrand(&prng))*sum,  sum1=0; uint16_t        ret=0;
		for(uint32_t      i=0;    i<O;    i++){   sum1+=out.out.data[i];    if(sum1>=ran){  ret=i;  break;  }       }
		curr=(curr+1)%C;	return	code[ret];
	}
	string	generate(string	inp,	uint32_t	n){
		vector<uint16_t>	v,s;	uint16_t	c;	v.resize(inp.size());
		v.resize(utf8_to_utf16((uint8_t*)inp.data(),inp.size(),v.data(),v.size()));
		if(!v.size())	return	"";
		for(uint32_t	i=0;	i<v.size()&&i<n;	i++){
			s.push_back(v[i]);	memset(vs,0,sizeof(float)*O);
			for(uint32_t	k=0;	k<s.size();	k++){
				uint32_t	l=1;
				while(l<=k&&s[k-l]==s[s.size()-l])	l++;
				vs[table[s[k]]]+=(expf(l)-1)/(s.size()-k);
			}
			c=sample(s.data()+s.size()-1,s.data());
		}
		while(s.size()<n){
			s.push_back(c);	memset(vs,0,sizeof(float)*O);
			for(uint32_t	k=0;	k<s.size();	k++){
				uint32_t	l=1;
				while(l<=k&&s[k-l]==s[s.size()-l])	l++;
				vs[table[s[k]]]+=(expf(l)-1)/(s.size()-k);
			}
			c=sample(s.data()+s.size()-1,s.data());
		}
		string	ret;	ret.resize(s.size()*4);
		ret.resize(utf16_to_utf8(s.data(),s.size(),(uint8_t*)ret.data(),ret.size()));
		return	ret;
	}
	float	probability(const	uint16_t	*x,	const	uint16_t	*p){
		uint32_t	para=p+C-1>=x?p+C-1-x:0;
		for(uint32_t	r=0;	r<E;	r++)	emb.data[r]=(_wyhash64(*x,r)&1)*2-1.0f;
		for(uint32_t	d=0;	d<D;	d++)	tra[d].fw(d?tra[d-1].out:emb,curr,para);
		layernorm<E,H>(tra[D-1].out);	out.fw(tra[D-1].out);
		softmax<O>(out.out);	
		curr=(curr+1)%C;	return	out.out.data[table[x[1]]];
	}
	float	evaluate(string	inp){
		vector<uint16_t>	v;	v.resize(inp.size());
		uint32_t	m=utf8_to_utf16((uint8_t*)inp.data(),inp.size(),v.data(),v.size());
		v.resize(m);	double	loss=0;
		for(uint32_t	i=0;	i+1<v.size();	i++)	loss-=logf(fmaxf(probability(v.data()+i,v.data()),FLT_MIN));
		return	m<2?0:loss/(m-1);
	}
};
#ifndef	WEIXIN
#include	"config"
int	main(int	ac,	char	**av){
	Denisovan<context,embed,depth,heads,voca>	model;
	string	model_file="model";	
	int	opt,	threads=1;
	while((opt=getopt(ac,	av,	"m:t:"))>=0){
		switch(opt){
		case	'm':	model_file=optarg;	break;
		case	't':	threads=atoi(optarg);	break;
		}
	}
	omp_set_num_threads(threads);
	if(!model.load(model_file.c_str())){	fprintf(stderr,"fail to load %s\n",model_file.c_str());	return	0;	}
	timeval	beg,end;	gettimeofday(&beg,NULL);
	cout<<model.generate(av[optind],context)<<'\n';
	gettimeofday(&end,NULL);
	cerr<<end.tv_sec-beg.tv_sec+1e-6*(end.tv_usec-beg.tv_usec)<<'\n';
	return	0;
}
#else
#include	"mongoose.h"
#include	"config"
using	namespace	std;
static struct mg_serve_http_opts s_http_server_opts;
Denisovan<context,embed,depth,heads,voca>	model;
unordered_map<string,string>	ma;
static void ev_handler(struct mg_connection *nc, int ev, void *p) {
	if (ev != MG_EV_HTTP_REQUEST)	return;
	struct http_message	*hm=(struct http_message *) p;
	const	char	*beg=strstr(hm->body.p,"<Content><![CDATA[");
	const	char	*end=strstr(hm->body.p,"]]></Content>");
	if(beg==NULL||end==NULL)	return;
	string	query;	for(const	char	*p=beg+strlen("<Content><![CDATA[");	p<end;	p++)	query.push_back(*p);
	beg=strstr(hm->body.p,"<FromUserName><![CDATA[");
	end=strstr(hm->body.p,"]]></FromUserName>");
	if(beg==NULL||end==NULL)	return;
	string	user;	for(const   char    *p=beg+strlen("<FromUserName><![CDATA[");    p<end;  p++)    user.push_back(*p);
	beg=strstr(hm->body.p,"<ToUserName><![CDATA[");
	end=strstr(hm->body.p,"]]></ToUserName>");
	if(beg==NULL||end==NULL)	return;
	string	myid;	for(const   char    *p=beg+strlen("<ToUserName><![CDATA[");    p<end;  p++)    myid.push_back(*p);
	beg=strstr(hm->body.p,"<MsgId>");
	end=strstr(hm->body.p,"</MsgId>");
	if(beg==NULL||end==NULL)	return;
	string	mid;	for(const   char    *p=beg+strlen("<MsgId>");    p<end;  p++)    mid.push_back(*p);
	string	ret;
	auto	mi=ma.find(mid);
	if(mi!=ma.end())	ret=mi->second;
	else{
		timeval	bt,et;	gettimeofday(&bt,NULL);
		ret=model.generate(query,672);
		ma[mid]=ret;
		gettimeofday(&et,NULL);	cerr<<et.tv_sec-bt.tv_sec+1e-6*(et.tv_usec-bt.tv_usec)<<'\n';
	}
	ostringstream	so;	so.precision(4);	so.setf(ios::fixed);
	so<<"<xml>\n";
  	so<<"<ToUserName><![CDATA["<<user<<"]]></ToUserName>\n";
  	so<<"<FromUserName><![CDATA["<<myid<<"]]></FromUserName>\n";
	so<<"<CreateTime>"<<time(NULL)<<"</CreateTime>\n";
  	so<<"<MsgType><![CDATA[text]]></MsgType>\n";
  	so<<"<Content><![CDATA["<<ret<<"]]></Content>\n";
	so<<"</xml>";
	mg_send_head(nc, 200, so.str().size(), "");
	mg_printf(nc,	"%s",	so.str().c_str());
	mg_serve_http(nc, (struct http_message *) p, s_http_server_opts);
}
int main(int	ac,	char	**av) {
  if(ac<2){	cerr<<"weixin port\n";	return	0;	}
  model.load("model");
  struct mg_mgr mgr;
  struct mg_connection *nc;
  mg_mgr_init(&mgr, NULL);
  nc = mg_bind(&mgr, av[1], ev_handler);
  mg_set_protocol_http_websocket(nc);
  for(;;)	mg_mgr_poll(&mgr, 15000);
  mg_mgr_free(&mgr);
  return 0;
}
#endif
